from langchain import PromptTemplate, OpenAI
from langchain.chains.llm import LLMChain
from langchain.prompts import PromptTemplate

#
llm = OpenAI(
    api_key="sk-VowKQBUMIkSND8WScNJtDLqf3FyqWHQ43LMVUXH1m6GZaopA",
    base_url="https://ai.nengyongai.cn/v1",
    temperature=0.8)


def test1():
    # 提示词
    prompt = PromptTemplate(
        template='给我讲一个笑话',
        input_variables=[]
    )
    # 格式化
    text = prompt.format()
    print(text)


def test2():
    # 提示词
    prompt = PromptTemplate(
        template='给我讲一个关于{content}的{adjective}笑话',
        input_variables=['adjective', 'content']
    )

    # 格式化
    text = prompt.format(adjective='有趣', content='鸡')
    print(text)


def test44():
    # 提示词
    prompt = PromptTemplate(
        template="我正在开一家新的{store_type}店铺，请帮我想出 5 个有创意的名字。",
        input_variables=["store_type"]
    )

    # 链
    chain = LLMChain(llm=llm, prompt=prompt)
    # 执行
    result = chain.run({"store_type": "咖啡馆"})  # 指定提示词变量值
    print(result)
